#!/usr/bin/env python3
import os
import click
import tomli
import sys
import pandas as pd
from tasks.model.task import TaskSpec
from tasks.task_gen import gen_task_spec
from tasks.server_helper import ServerHelper
from loguru import logger
from tasks.cfg_parser import parse_cfg
import asyncio
from tasks.worker.worker_manager import WorkerManager
import requests
from pprint import pprint

def discover_server_url():
    """Attempt to discover the task server URL.

    Tries different URLs in order and returns the first one that responds with 200 OK.
    Falls back to localhost:9010 if no server is found.
    """
    potential_urls = ["http://localhost:9010/", "http://server.tail13fe1.ts.net:9010/"]

    for url in potential_urls:
        try:
            response = requests.get(url, timeout=1)
            if response.status_code == 200:
                base_url = url.rstrip("/")
                logger.info(f"Discovered server at {base_url}")
                return base_url
        except requests.RequestException:
            continue

    # Default fallback
    default_url = "http://localhost:9010"
    logger.warning(f"Could not discover server, falling back to {default_url}")
    return default_url


@click.group()
def cli():
    pass

def _toml_file_to_task_spec(toml_file: str, commit_changes: bool = True) -> TaskSpec:
    cfg = tomli.load(open(toml_file, "rb"))
    root = cfg.get("__root__")
    if root is None:
        root = os.path.dirname(toml_file)
    elif not os.path.isabs(root):
        root = os.path.abspath(os.path.join(os.path.dirname(toml_file), root))
    cfg["__root__"] = root
    cfg = parse_cfg(cfg)
    pprint(cfg)
    _, task_specs, _ = gen_task_spec(cfg, None, commit_changes=commit_changes)
    return task_specs

@cli.command()
@click.argument("toml_file", type=click.Path(exists=True))
def parse(toml_file):
    """Parse a task from a TOML file."""
    task_spec = _toml_file_to_task_spec(toml_file, commit_changes=False)
    for k, v in task_spec.items():
        print(k)
        print(v.model_dump_json(indent=2))
        print()


@cli.command()
@click.argument("toml_file", type=click.Path(exists=True))
@click.argument("task", nargs=-1, required=False)
@click.option("--server-url", default=None, help="Server URL to push tasks to")
def run(toml_file, task, server_url):
    """Generate task specs from a TOML config file and push them to the server."""
    if not task:
        task = None

    # Discover server URL if not provided
    if server_url is None:
        server_url = discover_server_url()
    logger.info(f"Using server URL: {server_url}")

    try:
        task_specs = _toml_file_to_task_spec(toml_file, commit_changes=True)
        pprint(task_specs)
    except Exception as e:
        click.echo(f"Failed to load TOML: {e}", err=True)
        sys.exit(1)
    try:
        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        # Push all task specs to server
        click.echo("Pushing tasks to server...")
        for task_spec in task_specs.values():
            try:
                _, created = server.create_task_spec(task_spec)
                if created:
                    click.echo(f"Successfully pushed task: {task_spec.task}")
                else:
                    click.echo(f"Task already exists: {task_spec.task}")
            except Exception as e:
                click.echo(f"Failed to push task {task_spec.task}: {e}", err=True)
                logger.exception(e)
                break
        click.echo("\nTask push completed!")
    except Exception as e:
        click.echo(f"Error: {e}", err=True)
        logger.exception(e)
        sys.exit(2)


@cli.command()
@click.option("-n", "--num", default=10, help="Number of tasks to display")
@click.option(
    "--server-url",
    default=None,
    help="Server URL to fetch tasks from",
)
def list(num, server_url):
    """List recently added tasks in a tabular format."""
    try:
        # Discover server URL if not provided
        if server_url is None:
            server_url = discover_server_url()
        logger.info(f"Using server URL: {server_url}")

        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        # Fetch task specs
        task_specs = server.list_task_specs(limit=num)

        if not task_specs:
            click.echo("No tasks found.")
            return

        # Create DataFrame with relevant information
        data = []
        for spec in task_specs:
            data.append(
                {
                    "Task Name": spec.task,
                    "Status": spec.status,
                    "Task Hash": spec.task_hash,
                    "Commit": spec.commit[:7],  # Short hash for readability
                }
            )

        # Create and display DataFrame
        df = pd.DataFrame(data)
        with pd.option_context(
            "display.max_rows",
            None,
            "display.max_colwidth",
            None,
            "display.width",
            None,
            "display.max_columns",
            None,
        ):
            click.echo(df.to_string(index=False))

    except Exception as e:
        click.echo(f"Error listing tasks: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


@cli.command()
@click.argument("task_hash", nargs=-1, required=True)
@click.option(
    "--server-url",
    default=None,
    help="Server URL to delete task from",
)
def delete(task_hash, server_url):
    """Delete tasks and all their dependent tasks from the server.

    This will delete the task records from the database and all associated files.
    All tasks that depend on these tasks will also be deleted.
    """
    try:
        # Discover server URL if not provided
        if server_url is None:
            server_url = discover_server_url()
        logger.info(f"Using server URL: {server_url}")

        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        total_deleted = 0
        for hash_value in task_hash:
            try:
                # Perform delete operation
                click.echo(f"Deleting task {hash_value} and all dependent tasks...")
                deleted_tasks = server.delete_task(hash_value)

                click.echo(f"Successfully deleted {len(deleted_tasks)} task(s):")
                for deleted_hash in deleted_tasks:
                    click.echo(f"  - {deleted_hash}")

                total_deleted += len(deleted_tasks)
            except Exception as e:
                click.echo(f"Error deleting task {hash_value}: {e}", err=True)
                logger.exception(e)

        click.echo(f"\nTotal tasks deleted: {total_deleted}")

    except Exception as e:
        click.echo(f"Error in delete operation: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


@cli.command()
@click.argument("task_hash", nargs=-1, required=True)
@click.option(
    "--server-url",
    default=None,
    help="Server URL to reset task from",
)
def reset(task_hash, server_url):
    """Reset tasks and all their dependent tasks back to PENDING status.

    This will set the task status back to PENDING and delete any associated files.
    All tasks that depend on these tasks will also be reset.
    The task records and dependency records will not be removed.
    """
    try:
        # Discover server URL if not provided
        if server_url is None:
            server_url = discover_server_url()
        logger.info(f"Using server URL: {server_url}")

        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        total_reset = 0
        for hash_value in task_hash:
            try:
                # Perform reset operation
                click.echo(f"Resetting task {hash_value} and all dependent tasks...")
                reset_tasks = server.reset_task(hash_value)

                click.echo(f"Successfully reset {len(reset_tasks)} task(s):")
                for reset_hash in reset_tasks:
                    click.echo(f"  - {reset_hash}")

                total_reset += len(reset_tasks)
            except Exception as e:
                click.echo(f"Error resetting task {hash_value}: {e}", err=True)
                logger.exception(e)

        click.echo(f"\nTotal tasks reset: {total_reset}")

    except Exception as e:
        click.echo(f"Error in reset operation: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


@cli.command()
@click.argument("task_hash")
@click.option(
    "--server-url",
    default=None,
    help="Server URL to inspect task from",
)
@click.option(
    "--depth",
    default=5,
    help="Maximum depth of dependency tree to display",
)
def inspect(task_hash, server_url, depth):
    """Inspect a task and display its dependencies in a tree structure.

    This command displays the task and its dependencies in a hierarchical tree
    structure, using indentation to represent the dependency levels.
    """
    try:
        # Discover server URL if not provided
        if server_url is None:
            server_url = discover_server_url()
        logger.info(f"Using server URL: {server_url}")

        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        # Get the task spec
        try:
            task_spec = server.get_task_spec(task_hash)
            click.echo(f"Task: {task_spec.task} ({task_hash})")
            click.echo(f"Status: {task_spec.status}")
        except Exception as e:
            click.echo(f"Error getting task spec: {e}", err=True)
            # Continue anyway to show dependencies

        # Build and display dependency tree
        click.echo("\nDependency Tree:")
        _print_dependency_tree(server, task_hash, depth, indent="", visited=set())

    except Exception as e:
        click.echo(f"Error inspecting task: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


@cli.command()
@click.option(
    "--server-url",
    default=None,
    help="Server URL to fetch tasks from",
)
@click.option(
    "--workers",
    "-w",
    default=1,
    help="Number of worker processes to run",
)
@click.option(
    "--n-tasks",
    "-n",
    default=None,
    type=int,
    help="Number of tasks to process before exiting (None means run indefinitely)",
)
@click.option(
    "--cache-dir",
    default=None,
    help="Directory for task cache (defaults to ~/task_cache)",
)
def worker(server_url, workers, n_tasks, cache_dir):
    """Run a worker manager that processes tasks from the server.

    This command starts a worker manager that will continuously poll the server
    for pending tasks and execute them. The worker will run until it is interrupted
    or until it has processed the specified number of tasks (if n-tasks is provided).
    """
    try:
        # Discover server URL if not provided
        if server_url is None:
            server_url = discover_server_url()
        logger.info(f"Using server URL: {server_url}")

        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        # Initialize worker manager
        worker_manager = WorkerManager(server_helper=server, n_workers=workers, cache_dir=cache_dir)

        click.echo(f"Starting worker manager with {workers} workers...")
        if n_tasks:
            click.echo(f"Will process {n_tasks} tasks and then exit")
        else:
            click.echo("Will run until interrupted")

        # Run the worker manager
        asyncio.run(worker_manager.loop(n_tasks=n_tasks))

    except KeyboardInterrupt:
        click.echo("\nWorker manager interrupted, shutting down...")
    except Exception as e:
        click.echo(f"Error running worker manager: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


@cli.command()
@click.option(
    "--server-url",
    default=None,
    help="Server URL to prune tasks from",
)
@click.option(
    "--dry-run",
    is_flag=True,
    help="Show what would be pruned without actually deleting",
)
def prune(server_url, dry_run):
    """Remove unused task data where files exist but task is already deleted.
    
    This command compares the list of files in storage with the list of tasks
    in the database and removes any files that don't correspond to existing tasks.
    """
    try:
        # Discover server URL if not provided
        if server_url is None:
            server_url = discover_server_url()
        logger.info(f"Using server URL: {server_url}")

        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        # Get all files
        files = set((s[:-1] for s in server.list_files()))
        
        # Get all task hashes
        all_tasks = set()
        chunk = 10
        for i in range(0, 99999999, chunk):
            tasks = server.list_task_specs(offset=i, limit=chunk)
            for task in tasks:
                all_tasks.add(task.task_hash)
            if len(tasks) < chunk:
                break

        # Find files to prune
        to_be_pruned = files - all_tasks

        if not to_be_pruned:
            click.echo("No unused task data found.")
            return

        click.echo(f"Found {len(to_be_pruned)} unused task data entries:")
        for task_hash in to_be_pruned:
            click.echo(f"  - {task_hash}")

        if dry_run:
            click.echo("\nDry run - no files were deleted")
            return

        # Delete the files
        total_pruned = 0
        for task_hash in to_be_pruned:
            try:
                server._delete_task_files(task_hash)
                total_pruned += 1
            except Exception as e:
                click.echo(f"Error pruning task {task_hash}: {e}", err=True)
                logger.exception(e)

        click.echo(f"\nSuccessfully pruned {total_pruned} unused task data entries")

    except Exception as e:
        click.echo(f"Error in prune operation: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


def _print_dependency_tree(
    server, task_hash, max_depth, indent="", visited=None, current_depth=0
):
    """
    Recursively print a tree of task dependencies.

    Args:
        server: ServerHelper instance
        task_hash: Current task hash to inspect
        max_depth: Maximum depth to traverse
        indent: Current indentation string
        visited: Set of already visited task hashes to prevent cycles
        current_depth: Current depth level
    """
    if visited is None:
        visited = set()

    # Prevent cycles and respect max depth
    if task_hash in visited or current_depth >= max_depth:
        return

    visited.add(task_hash)

    # Get direct dependencies
    try:
        dependencies = server.get_dependencies(task_hash)

        # If no dependencies, print a leaf node
        if not dependencies:
            click.echo(f"{indent}└── {task_hash} (no dependencies)")
            return

        # Try to get task name for each dependency
        for i, dep_hash in enumerate(dependencies):
            is_last = i == len(dependencies) - 1

            # Choose the appropriate branch symbol
            branch = "└── " if is_last else "├── "

            # Get task name if possible
            try:
                dep_task = server.get_task_spec(dep_hash)
                task_name = dep_task.task
                task_status = dep_task.status.value
                click.echo(f"{indent}{branch}{dep_hash} ({task_name}, {task_status})")
            except Exception:
                # If we can't get the task details, just show the hash
                click.echo(f"{indent}{branch}{dep_hash}")

            # Prepare the next level indentation
            next_indent = indent + ("    " if is_last else "│   ")

            # Recursively print dependencies of this dependency
            _print_dependency_tree(
                server, dep_hash, max_depth, next_indent, visited, current_depth + 1
            )

    except Exception as e:
        click.echo(f"{indent}└── Error getting dependencies: {e}", err=True)


@cli.command()
@click.argument("local_file", type=click.Path(exists=True))
@click.option("--remote-path", default=None, help="Remote path on server (defaults to filename only)")
@click.option(
    "--server-url",
    default=None,
    help="Server URL to upload file to",
)
def upload(local_file, remote_path, server_url):
    """Upload a local file to the server.
    
    If no remote path is specified, the file will be uploaded using just the filename
    (without the local directory path) as the remote path.
    
    Examples:
        pytask upload /path/to/myfile.txt
        pytask upload /path/to/myfile.txt --remote-path data/myfile.txt
    """
    try:
        # Discover server URL if not provided
        if server_url is None:
            server_url = discover_server_url()
        logger.info(f"Using server URL: {server_url}")

        # If no remote path specified, use just the filename
        if remote_path is None:
            remote_path = os.path.basename(local_file)
            click.echo(f"No remote path specified, using filename: {remote_path}")

        # Initialize server helper
        server = ServerHelper(base_url=server_url)

        # Upload the file
        click.echo(f"Uploading {local_file} to {remote_path}...")
        server.upload_file(local_file, remote_path)
        
        click.echo(f"Successfully uploaded {local_file} to {remote_path}")

    except FileNotFoundError as e:
        click.echo(f"File not found: {e}", err=True)
        sys.exit(1)
    except IsADirectoryError as e:
        click.echo(f"Cannot upload directory: {e}", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error uploading file: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


if __name__ == "__main__":
    cli()
